{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from cholespy import CholeskySolverF, MatrixType\n",
    "from utils import get_icosphere, get_coo_arrays\n",
    "import igl\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "import seaborn as sns\n",
    "import time\n",
    "import numpy as np\n",
    "import tqdm\n",
    "\n",
    "import sksparse.cholmod as cholmod\n",
    "import scipy.sparse as sp\n",
    "import cupy as cp\n",
    "import cupyx.scipy.sparse as cps\n",
    "import cupyx.scipy.sparse.linalg as cpsl\n",
    "from cupy import cusparse\n",
    "from torch.utils.dlpack import to_dlpack\n",
    "from torch.utils.dlpack import from_dlpack\n",
    "\n",
    "matplotlib.rcParams['font.size'] = 18\n",
    "matplotlib.rcParams['text.usetex'] = True\n",
    "matplotlib.rcParams['text.latex.preamble'] = r\"\"\"\\usepackage{libertine}\n",
    "\\usepackage{amsmath}\"\"\"\n",
    "matplotlib.rcParams['pdf.fonttype'] = 42\n",
    "matplotlib.rcParams['ps.fonttype'] = 42\n",
    "sns.set(font_scale=1.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def torch_to_cupy(x):\n",
    "    return cp.from_dlpack(to_dlpack(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cupy_to_torch(x):\n",
    "    return from_dlpack(x.toDlpack())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare(A, transpose, blocking=True, level_info=True):\n",
    "    import cupy as _cupy\n",
    "    from cupy.cuda import device as _device\n",
    "    from cupy_backends.cuda.libs import cusparse as _cusparse\n",
    "    import cupyx.scipy.sparse\n",
    "\n",
    "    policy = _cusparse.CUSPARSE_SOLVE_POLICY_USE_LEVEL if level_info \\\n",
    "        else _cusparse.CUSPARSE_SOLVE_POLICY_NO_LEVEL\n",
    "    algo = 1 if blocking else 0\n",
    "\n",
    "    transa = _cusparse.CUSPARSE_OPERATION_TRANSPOSE if transpose \\\n",
    "        else _cusparse.CUSPARSE_OPERATION_NON_TRANSPOSE\n",
    "    transb = _cusparse.CUSPARSE_OPERATION_TRANSPOSE\n",
    "    fill_mode = _cusparse.CUSPARSE_FILL_MODE_LOWER\n",
    "\n",
    "    if cupyx.scipy.sparse.isspmatrix_csc(A):\n",
    "        A = A.T\n",
    "        transa = 1 - transa\n",
    "        fill_mode = 1 - fill_mode\n",
    "\n",
    "    assert cupyx.scipy.sparse.isspmatrix_csr(A)\n",
    "\n",
    "    handle = _device.get_cusparse_handle()\n",
    "    info = _cusparse.createCsrsm2Info()\n",
    "    m = A.shape[0]\n",
    "    alpha = np.array(1, dtype=np.float32)\n",
    "    desc = _cusparse.createMatDescr()\n",
    "    _cusparse.setMatType(desc, _cusparse.CUSPARSE_MATRIX_TYPE_GENERAL)\n",
    "    _cusparse.setMatIndexBase(desc, _cusparse.CUSPARSE_INDEX_BASE_ZERO)\n",
    "    _cusparse.setMatFillMode(desc, fill_mode)\n",
    "    _cusparse.setMatDiagType(desc, _cusparse.CUSPARSE_DIAG_TYPE_NON_UNIT)\n",
    "\n",
    "    nrhs = 128\n",
    "    ldb = nrhs\n",
    "\n",
    "    ws_size = _cusparse.scsrsm2_bufferSizeExt(\n",
    "        handle, algo, transa, transb, m, nrhs, A.nnz, alpha.ctypes.data,\n",
    "        desc, A.data.data.ptr, A.indptr.data.ptr, A.indices.data.ptr,\n",
    "        0, ldb, info, policy)\n",
    "\n",
    "    ws = _cupy.empty((ws_size,), dtype=np.int8)\n",
    "\n",
    "    _cusparse.scsrsm2_analysis(\n",
    "         handle, algo, transa, transb, m, nrhs, A.nnz, alpha.ctypes.data,\n",
    "         desc, A.data.data.ptr, A.indptr.data.ptr,\n",
    "         A.indices.data.ptr, 0, ldb, info, policy, ws.data.ptr)\n",
    "\n",
    "    def solver(b):\n",
    "        _cusparse.scsrsm2_solve(\n",
    "            handle, algo, transa, transb, m, nrhs, A.nnz, alpha.ctypes.data,\n",
    "            desc, A.data.data.ptr, A.indptr.data.ptr, A.indices.data.ptr,\n",
    "            b.data.ptr, ldb, info, policy, ws.data.ptr)\n",
    "\n",
    "    return solver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Cholesky:\n",
    "    \"\"\"\n",
    "    Cholesky decomposition of a matrix\n",
    "    \"\"\"\n",
    "    def __init__(self, n_verts, M_cpu):\n",
    "        \"\"\"\n",
    "        M: pytorch sparse tensor, the matrix to decompose. Assumed to be spd\n",
    "        \"\"\"\n",
    "\n",
    "        factor = cholmod.cholesky(M_cpu, ordering_method='nesdis', mode='simplicial')\n",
    "        L, P = factor.L(), factor.P()\n",
    "        # Invert the permutation\n",
    "        Pi = np.argsort(P).astype(np.int32)\n",
    "\n",
    "        # Transfer to GPU as cupy arrays\n",
    "        self.L = cps.csc_matrix(L.astype(np.float32))\n",
    "        self.U = self.L.T\n",
    "        self.P = cp.array(P)\n",
    "        self.Pi = cp.array(Pi)\n",
    "        \n",
    "        self.solver_1 = prepare(self.L, False, False, True)\n",
    "        self.solver_2 = prepare(self.L, True, False, True)\n",
    "\n",
    "    def solve(self, b):\n",
    "        tmp = torch_to_cupy(b)[self.P]\n",
    "        self.solver_1(tmp)\n",
    "        self.solver_2(tmp)\n",
    "        sol = cupy_to_torch(tmp[self.Pi])\n",
    "        \n",
    "        return sol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lambda_ = 2.0\n",
    "n_verts, f = get_icosphere(6)\n",
    "n_faces = len(f)\n",
    "n_runs = np.array([20, 100])\n",
    "n_rhs = 128\n",
    "timings_cupy = np.zeros(2, dtype=float)\n",
    "timings_ours_cuda = np.zeros(2, dtype=float)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Baseline GPU timing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in tqdm.trange(n_runs[0]):\n",
    "    # change the matrix, but don't time this part\n",
    "    values, idx = get_coo_arrays(n_verts, f, i+1)\n",
    "    M_cpu = sp.coo_matrix((values, idx))\n",
    "    \n",
    "    start = torch.cuda.Event(enable_timing=True)\n",
    "    end = torch.cuda.Event(enable_timing=True)\n",
    "       \n",
    "    start.record()\n",
    "    \n",
    "    solver_cupy = Cholesky(n_verts, M_cpu) \n",
    "    \n",
    "    end.record()\n",
    "    torch.cuda.synchronize()\n",
    "    \n",
    "    timings_cupy[0] += start.elapsed_time(end)\n",
    "\n",
    "for i in tqdm.trange(n_runs[1]):\n",
    "    \n",
    "    b_numpy = np.random.uniform(size=(n_verts, n_rhs))\n",
    "    b_torch = torch.tensor(b_numpy, device='cuda', dtype=torch.float32)\n",
    "    \n",
    "    start = torch.cuda.Event(enable_timing=True)\n",
    "    end = torch.cuda.Event(enable_timing=True)\n",
    "       \n",
    "    start.record()\n",
    "    \n",
    "    x_cusparse = solver_cupy.solve(b_torch)\n",
    "    \n",
    "    end.record()\n",
    "    torch.cuda.synchronize()\n",
    "    \n",
    "    timings_cupy[1] += start.elapsed_time(end)\n",
    "\n",
    "timings_cupy /= n_runs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Our GPU timing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in tqdm.trange(n_runs[0]):\n",
    "    \n",
    "    values, idx = get_coo_arrays(n_verts, f, i+1)\n",
    "    ii = torch.tensor(idx[0], device='cuda')\n",
    "    jj = torch.tensor(idx[1], device='cuda')\n",
    "    x = torch.tensor(values, device='cuda')\n",
    "    \n",
    "    start = torch.cuda.Event(enable_timing=True)\n",
    "    end = torch.cuda.Event(enable_timing=True)\n",
    "       \n",
    "    start.record()\n",
    "    solver_ours = CholeskySolverF(n_verts, ii, jj, x, MatrixType.COO)\n",
    "    \n",
    "    end.record()\n",
    "    torch.cuda.synchronize()\n",
    "    \n",
    "    timings_ours_cuda[0] += start.elapsed_time(end)\n",
    "\n",
    "for i in tqdm.trange(n_runs[1]):\n",
    "    \n",
    "    b_numpy = np.random.uniform(size=(n_verts, n_rhs))\n",
    "    b_torch = torch.tensor(b_numpy, device='cuda', dtype=torch.float32)\n",
    "    \n",
    "    start = torch.cuda.Event(enable_timing=True)\n",
    "    end = torch.cuda.Event(enable_timing=True)\n",
    "       \n",
    "    start.record()\n",
    "    \n",
    "    x_torch = torch.zeros_like(b_torch)\n",
    "    solver_ours.solve(b_torch, x_torch)\n",
    "    \n",
    "    end.record()\n",
    "    torch.cuda.synchronize()\n",
    "    \n",
    "    timings_ours_cuda[1] += start.elapsed_time(end)\n",
    "\n",
    "\n",
    "timings_ours_cuda /= n_runs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Baseline CPU timing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "timings_sksparse = np.zeros(2, dtype=float)\n",
    "timings_ours_cpu = np.zeros(2, dtype=float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in tqdm.trange(n_runs[0]):\n",
    "    # change the matrix, but don't time this part\n",
    "    values, idx = get_coo_arrays(n_verts, f, i+1)\n",
    "    M_cpu = sp.coo_matrix((values, idx))\n",
    "    \n",
    "    start = time.perf_counter()\n",
    "    \n",
    "    solver_sksparse = cholmod.cholesky(M_cpu, ordering_method='nesdis', mode='simplicial')\n",
    "\n",
    "    end = time.perf_counter()\n",
    "    \n",
    "    timings_sksparse[0] += end - start\n",
    "\n",
    "for i in tqdm.trange(n_runs[1]):\n",
    "    \n",
    "    b_numpy = np.random.uniform(size=(n_verts, n_rhs))\n",
    "    \n",
    "    start = time.perf_counter()\n",
    "    \n",
    "    x_sksparse = solver_sksparse(b_numpy)\n",
    "    \n",
    "    end = time.perf_counter()\n",
    "    \n",
    "    timings_sksparse[1] += end - start\n",
    "\n",
    "timings_sksparse /= n_runs / 1000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Our CPU timing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for i in tqdm.trange(n_runs[0]):\n",
    "    \n",
    "    values, idx = get_coo_arrays(n_verts, f, i+1)\n",
    "    \n",
    "    start = time.perf_counter()\n",
    "        \n",
    "    solver_ours = CholeskySolverF(n_verts, idx[0], idx[1], values, MatrixType.COO)\n",
    "    \n",
    "    end = time.perf_counter()\n",
    "    \n",
    "    timings_ours_cpu[0] += end - start\n",
    "\n",
    "for i in tqdm.trange(n_runs[1]):\n",
    "    \n",
    "    b_numpy = np.random.uniform(size=(n_verts, n_rhs)).astype(np.float32)\n",
    "    x_numpy = np.zeros_like(b_numpy)\n",
    "    \n",
    "    start = time.perf_counter()\n",
    "    \n",
    "    solver_ours.solve(b_numpy, x_numpy)\n",
    "    \n",
    "    end = time.perf_counter()\n",
    "    \n",
    "    timings_ours_cpu[1] += end - start\n",
    "\n",
    "\n",
    "timings_ours_cpu /= n_runs / 1000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "titles = [\"CUDA Version\", \"CPU Version\"]\n",
    "steps = ['Analysis', 'Solve']\n",
    "labels = [[\"CuSparse\", \"Ours (CUDA)\"], [\"Scikit-Sparse\", \"Ours (CPU)\"]]\n",
    "timings_ours = [timings_ours_cuda, timings_ours_cpu]\n",
    "timings_ref = [timings_cupy, timings_sksparse]\n",
    "\n",
    "x = np.arange(2)  # the label locations\n",
    "width = 0.35  # the width of the bars\n",
    "\n",
    "aspect = 4/3\n",
    "nrows = 1\n",
    "ncols = 2\n",
    "\n",
    "base_size = 10\n",
    "w = ncols * base_size\n",
    "h = nrows * base_size / aspect\n",
    "\n",
    "fontsize = 24\n",
    "\n",
    "fig, ax = plt.subplots(nrows=nrows, ncols=ncols, figsize=(w,h))\n",
    "\n",
    "for i, (names, t_ref, t_ours) in enumerate(zip(labels, timings_ref, timings_ours)):\n",
    "    rects1 = ax[i].bar(x - width/2, t_ref, width, label=names[0])\n",
    "    rects2 = ax[i].bar(x + width/2, t_ours, width, label=names[1])\n",
    "    \n",
    "    if i == 0:\n",
    "        ax[i].set_ylabel('Elapsed time (ms)', fontsize=fontsize)\n",
    "    \n",
    "    ax[i].set_xticks(x, steps, fontsize=18)\n",
    "    ax[i].legend()\n",
    "\n",
    "    bar_label_1 = [f\"{time:.2f}\" for time in t_ref]\n",
    "    bar_label_2 = [f\"{t_ours[i]:.2f} (x{t_ours[i]/t_ref[i]:.2f})\" for i in range(len(t_ref))]\n",
    "\n",
    "    ax[i].bar_label(rects1, labels=bar_label_1, fontsize=15)\n",
    "    ax[i].bar_label(rects2, labels=bar_label_2, fontsize=15)\n",
    "    \n",
    "    ax[i].set_title(titles[i], y=-0.2, fontsize=fontsize)\n",
    "    ax[i].set_ylim(5, 400)\n",
    "\n",
    "    plt.setp(ax[i], yscale='log')\n",
    "\n",
    "plt.savefig(\"benchmark.jpg\", format='jpg', dpi=500, bbox_inches='tight', pad_inches=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
